import os
import numpy as np
import matplotlib.pyplot as plt


def read_data_from_txt(file_path):
    y_coords = []
    confidences = []
    with open(file_path, "r") as file:
        for line in file:
            parts = line.strip().split()
            if len(parts) >= 6:
                y_center = float(parts[3])
                confidence = float(parts[5])
                if y_center < 0.05:
                    y_coords.append(y_center)
                    confidences.append(confidence)
    return y_coords, confidences


def collect_data_from_folder(folder_path):
    indexi = 0
    min_i = 600
    max_i = 800
    all_y_coords = []
    all_confidences = []
    for file_name in os.listdir(folder_path):
        if file_name.endswith(".txt"):
            indexi = indexi + 1
            if indexi > min_i and indexi < max_i:
                file_path = os.path.join(folder_path, file_name)
                y_coords, confidences = read_data_from_txt(file_path)
                all_y_coords.extend(y_coords)
                all_confidences.extend(confidences)
    return all_y_coords, all_confidences


def calculate_average_confidence(y_coords, confidences, bins):
    bin_indices = np.digitize(y_coords, bins)
    bin_means = [
        np.mean(
            [confidences[i] for i in range(len(confidences)) if bin_indices[i] == j]
        )
        for j in range(1, len(bins))
    ]
    return bin_means


# 文件夹路径
model1_folder = "test_enhanced/labels"
model2_folder = "test_og/labels"

# 区间划分
bins = np.linspace(0, 0.05, 6)  # 将纵坐标小于0.05的部分划分为5个区间

# 收集数据
model1_y_coords, model1_confidences = collect_data_from_folder(model1_folder)
model2_y_coords, model2_confidences = collect_data_from_folder(model2_folder)

# 计算每个区间的平均信度
model1_bin_means = calculate_average_confidence(
    model1_y_coords, model1_confidences, bins
)
model2_bin_means = calculate_average_confidence(
    model2_y_coords, model2_confidences, bins
)

# 绘制立方图
plt.figure(figsize=(10, 6))
bin_centers = (bins[:-1] + bins[1:]) / 2
width = (bins[1] - bins[0]) * 0.4

plt.bar(
    bin_centers - width / 2,
    model1_bin_means,
    width=width,
    label="TCone-YOLO",
    alpha=0.7,
)
plt.bar(
    bin_centers + width / 2, model2_bin_means, width=width, label="YOLOv5s", alpha=0.7
)

# 添加区间标签
interval_labels = [f"{bins[i]:.3f} - {bins[i+1]:.3f}" for i in range(len(bins) - 1)]
plt.xticks(bin_centers, interval_labels)

plt.xlabel("Y Coordinate Interval (0 to 0.05) (normalized)", fontsize=14)
plt.ylabel("Average Confidence", fontsize=14)
plt.legend(fontsize=12)
plt.grid(True, linestyle="--", alpha=0.6)
plt.tight_layout()

# 保存图像为高分辨率的png文件
plt.savefig("confidence_distribution.png", dpi=1000)

# 显示图像
plt.show()
